package org.finesys.chat.core.base.embedding.factory;


import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.finesys.ai.entity.AigcKnowledge;
import org.finesys.ai.entity.AigcModel;
import org.finesys.chat.core.base.client.ChatClient;
import org.finesys.chat.core.base.document.splitter.DocumentSplitter;
import org.finesys.chat.core.base.document.splitter.DocumentSplitters;
import org.finesys.chat.core.base.embedding.EmbeddingStore;
import org.finesys.chat.core.base.factory.ChatModelFactory;
import org.finesys.chat.core.base.segment.TextSegment;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

@Slf4j
@Component
@RequiredArgsConstructor
public class EmbeddingProvider {

    private final EmbeddingStoreFactory embeddingStoreFactory;
    private final KnowledgeStoreFactory knowledgeStoreFactory;
    private final ChatModelFactory chatModelFactory;

    /**
     * 根据知识库ID获取嵌入存储对象
     *
     * @param knowledgeId 知识库ID
     * @return 嵌入存储对象
     * @throws IllegalArgumentException 如果知识库不存在或未配置嵌入存储，则抛出此异常
     */
    public EmbeddingStore<TextSegment> getEmbeddingStore(String knowledgeId) {
        if (!knowledgeStoreFactory.contains(knowledgeId)) {
            throw new IllegalArgumentException(String.format("知识库不存在，请先创建知识库：[%s]", knowledgeId));
        }
        AigcKnowledge knowledge = knowledgeStoreFactory.get(knowledgeId);
        //获取嵌入存储的标识符
        String embedStoreId = knowledge.getEmbedStoreId();
        if (!StringUtils.hasText(embedStoreId)) {
            throw new IllegalArgumentException(String.format("知识库未配置嵌入存储：[%s]", knowledgeId));
        }
        //获取模型信息
        String embeddingModelId = knowledge.getEmbedModelId();
        if (!StringUtils.hasText(embeddingModelId)) {
            throw new IllegalArgumentException(String.format("知识库未配置模型：[%s]", knowledgeId));
        }

        return embeddingStoreFactory.get(embedStoreId);
    }

    /**
     * 根据知识库ID获取聊天客户端
     *
     * @param knowledgeId 知识库ID
     * @return ChatClient 聊天客户端对象
     * @throws IllegalArgumentException 如果知识库不存在、未配置嵌入存储或未配置模型，则抛出此异常
     */
    public ChatClient getChatClient(String knowledgeId) {
        if (!knowledgeStoreFactory.contains(knowledgeId)) {
            throw new IllegalArgumentException(String.format("知识库不存在，请先创建知识库：[%s]", knowledgeId));
        }
        AigcKnowledge knowledge = knowledgeStoreFactory.get(knowledgeId);
        //获取嵌入存储的标识符
        String embedStoreId = knowledge.getEmbedStoreId();
        if (!StringUtils.hasText(embedStoreId)) {
            throw new IllegalArgumentException(String.format("知识库未配置嵌入存储：[%s]", knowledgeId));
        }
        //获取模型信息
        String embeddingModelId = knowledge.getEmbedModelId();
        if (!StringUtils.hasText(embeddingModelId)) {
            throw new IllegalArgumentException(String.format("知识库未配置模型：[%s]", knowledgeId));
        }
        return chatModelFactory.getChatClient(embeddingModelId);
    }

    /**
     * 根据知识库ID获取AIGC模型
     *
     * @param knowledgeId 知识库ID
     * @return AIGC模型
     * @throws IllegalArgumentException 如果知识库不存在，或者未配置嵌入存储，或者未配置模型，将抛出此异常
     */
    public AigcModel getAigcModel(String knowledgeId) {
        if (!knowledgeStoreFactory.contains(knowledgeId)) {
            throw new IllegalArgumentException(String.format("知识库不存在，请先创建知识库：[%s]", knowledgeId));
        }
        AigcKnowledge knowledge = knowledgeStoreFactory.get(knowledgeId);
        //获取嵌入存储的标识符
        String embedStoreId = knowledge.getEmbedStoreId();
        if (!StringUtils.hasText(embedStoreId)) {
            throw new IllegalArgumentException(String.format("知识库未配置嵌入存储：[%s]", knowledgeId));
        }
        //获取模型信息
        String embeddingModelId = knowledge.getEmbedModelId();
        if (!StringUtils.hasText(embeddingModelId)) {
            throw new IllegalArgumentException(String.format("知识库未配置模型：[%s]", knowledgeId));
        }
        return chatModelFactory.getModel(embeddingModelId);
    }


    public static DocumentSplitter splitter() {
        return DocumentSplitters.recursive(300, 20);
    }
}
